package xormplus

import (
	"errors"
	"fmt"
	"gitee.com/hongzhaomin/hzm-common-go/assert"
	"gitee.com/hongzhaomin/hzm-common-go/concurrent"
	"gitee.com/hongzhaomin/hzm-common-go/strutil"
	"sync/atomic"
	"xorm.io/xorm"
)

type DataSourceManager struct {
	primarySource *xorm.Engine // 主数据源
	engineMap     *concurrent.Map[string, *xorm.Engine]
}

func (dsm *DataSourceManager) Regis(dataSourceId string, engine *xorm.Engine) {
	assert.IsTrue(strutil.IsNotBlank(dataSourceId), "dataSourceId must not blank")
	assert.NonNil(engine, "engine must not nil")
	_, ok := dsm.engineMap.Load(dataSourceId)
	assert.IsFalse(ok, fmt.Sprintf("engine [%s] already registered", dataSourceId))
	dsm.engineMap.Store(dataSourceId, engine)
}

func (dsm *DataSourceManager) SetPrimary(dataSourceId string, engine *xorm.Engine) {
	assert.IsTrue(strutil.IsNotBlank(dataSourceId), "dataSourceId must not blank")
	assert.NonNil(engine, "engine must not nil")
	dsm.primarySource = engine
	dsm.engineMap.Store(dataSourceId, engine)
}

func (dsm *DataSourceManager) GetById(dataSourceId string) *xorm.Engine {
	// 数据源id为空
	if strutil.IsBlank(dataSourceId) {
		// 主数据源存在，返回主数据源
		if dsm.primarySource != nil {
			return dsm.primarySource
		}
		// 主数据源不存在，但是总数据源就只有一个，则返回这唯一的数据源
		if dsm.engineMap.Size() == 1 {
			// 如果 dataSourceId 为空，且 engine 只有一个，直接取出它即可
			var engine *xorm.Engine
			dsm.engineMap.Range(func(key string, value *xorm.Engine) (shouldContinue bool) {
				engine = value
				return false
			})
			return engine
		}
		panic(errors.New("DataSource not found, please regis it"))
	}
	engine, ok := dsm.engineMap.Load(dataSourceId)
	if !ok {
		panic(errors.New(fmt.Sprintf("DataSource [%s] not found, please regis it", dataSourceId)))
	}
	return engine
}

var dataSourceManager atomic.Pointer[DataSourceManager]

func init() {
	dataSourceManager.Store(&DataSourceManager{
		engineMap: concurrent.NewMap[string, *xorm.Engine](),
	})
}

func defaultDataSourceManager() *DataSourceManager {
	return dataSourceManager.Load()
}

func RegisDataSources(dataSourceId string, engine *xorm.Engine) {
	defaultDataSourceManager().Regis(dataSourceId, engine)
}

func SetPrimary(dataSourceId string, engine *xorm.Engine) {
	defaultDataSourceManager().SetPrimary(dataSourceId, engine)
}

func GetDataSourceById(dataSourceId string) *xorm.Engine {
	return defaultDataSourceManager().GetById(dataSourceId)
}

func GetDataSource() *xorm.Engine {
	return GetDataSourceById(strutil.Empty)
}
